package net.sansi.v3correctserver.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import javax.annotation.Nonnull;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;

/**
 * Created by Fang
 * create time : 2024/5/8
 * IDE : IntelliJ IDEA
 * <p>
 * 异步配置
 */
@Configuration
@EnableAsync
public class AsyncConfiguration implements AsyncConfigurer {

    private static final Logger log = LoggerFactory.getLogger(AsyncConfiguration.class);

    @Override
    public Executor getAsyncExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(20);
        executor.setMaxPoolSize(100);
        executor.setQueueCapacity(100000);
        executor.setThreadNamePrefix("Async-");
        executor.setKeepAliveSeconds(60);
        executor.setTaskDecorator(new RequestContextTaskDecorator());
        // 添加错误异常
        executor.setRejectedExecutionHandler((r, executor1) -> {
            throw new RejectedExecutionException("ThreadPool is full, cannot execute the task: " + r.toString());
        });
        executor.initialize();
        return executor;
    }

    /**
     * 异步请求上下文
     */
    static class RequestContextTaskDecorator implements TaskDecorator {
        @Override
        @Nonnull
        public Runnable decorate(@Nonnull Runnable runnable) {
            RequestAttributes reqAttr = RequestContextHolder.getRequestAttributes();
            if (reqAttr != null) {
                return () -> {
                    try {
                        RequestContextHolder.setRequestAttributes(reqAttr);
                        runnable.run();
                    } finally {
                        RequestContextHolder.resetRequestAttributes();
                    }
                };
            } else {
                return runnable;
            }
        }
    }
}
